﻿#ifndef IPC2_H
#define IPC2_H

#include "SharedMemory.h"
#include <memory>
#include <cstdint>
#include <functional>
#include <thread>
#include <vector>
#include <atomic>
#include <cstring>

enum class ERR_IPC
{
	E_NO_ERROR = 0,
	E_NO_ENOUGH_MEM,
	E_WAIT_READ,
	E_NOT_ATTACHED,
};

struct alignas(4) MemBlockHeader
{
	uint16_t subscriberCnt{0};
	uint16_t readCnt{0};
	uint16_t packetId{0};
	uint16_t reserved{0};
};

class MemBlock final
{
private:
	std::shared_ptr<SharedMemory> SharedMem{nullptr};

public:
	explicit MemBlock(const std::string &name, SharedMemory::AccessMode mode, size_t size = 10240)
	{
		SharedMem = std::make_shared<SharedMemory>(name);

		if (SharedMem->Attach(mode))
		{
			auto pMem = SharedMem->Data();
			auto memSize = SharedMem->DataLen();
		}
		else if (SharedMem->Create(size))
		{
			MemBlockHeader *pHeader = (MemBlockHeader *)SharedMem->Data();
			pHeader->subscriberCnt = 0;
			pHeader->readCnt = 0;
		}
		else
		{
			// can't create and attach to shared memory
		}
	}

	size_t Length() const
	{
		if (SharedMem && SharedMem->IsAttached())
		{
			constexpr const auto headerLen = sizeof(MemBlockHeader);
			return SharedMem->DataLen() - headerLen;
		}
		return 0;
	}

	ERR_IPC Write(void *data, size_t size)
	{
		if (!SharedMem->IsAttached())
		{
			return ERR_IPC::E_NOT_ATTACHED;
		}
		constexpr const auto headerLen = sizeof(MemBlockHeader);
		auto pBegin = SharedMem->Data() + headerLen;
		auto len = SharedMem->DataLen() - headerLen;
		if (len >= size)
		{
			SharedMem->Lock();
			MemBlockHeader *pHeader = (MemBlockHeader *)SharedMem->Data();
			if (pHeader->subscriberCnt == pHeader->readCnt ||
				pHeader->packetId == 0)
			{
				memcpy(pBegin, data, size);
				if (pHeader->packetId == 0xFFFF)
				{
					pHeader->packetId = 0;
				}
				pHeader->packetId++;
				SharedMem->Unlock();
				return ERR_IPC::E_NO_ERROR;
			}
			else
			{
				SharedMem->Unlock();
				return ERR_IPC::E_WAIT_READ;
			}
		}
		return ERR_IPC::E_NO_ENOUGH_MEM;
	}

	ERR_IPC Read(void *data, size_t size, uint16_t &packetId)
	{
		if (!SharedMem->IsAttached())
		{
			return ERR_IPC::E_NOT_ATTACHED;
		}
		constexpr const auto headerLen = sizeof(MemBlockHeader);
		auto pBegin = SharedMem->Data() + headerLen;
		auto len = SharedMem->DataLen() - headerLen;
		if (len >= size)
		{
			SharedMem->Lock();

			MemBlockHeader *pHeader = (MemBlockHeader *)SharedMem->Data();
			if (packetId != pHeader->packetId)
			{
				memcpy(data, pBegin, size);
				pHeader->readCnt++;
				packetId = pHeader->packetId;
			}

			SharedMem->Unlock();
			return ERR_IPC::E_NO_ERROR;
		}
		return ERR_IPC::E_NO_ENOUGH_MEM;
	}

	bool IsAttached() const
	{
		if (SharedMem && SharedMem->IsAttached())
		{
			return true;
		}
		return false;
	}

	void IncreaseSubscriber()
	{
		SharedMem->Lock();
		MemBlockHeader *pHeader = (MemBlockHeader *)SharedMem->Data();
		++pHeader->subscriberCnt;
		SharedMem->Unlock();
	}

	void DecreaseSubscriber()
	{
		SharedMem->Lock();
		MemBlockHeader *pHeader = (MemBlockHeader *)SharedMem->Data();
		--pHeader->subscriberCnt;
		SharedMem->Unlock();
	}
};

class Publisher final
{
public:
	explicit Publisher(const std::string &topic) : topic(topic)
	{
	}
	~Publisher() = default;

	/**
	 * @brief public data
	 * @return false:Not enough memory, or some subscribers have unread data
	 */
	ERR_IPC Publish(void *data, size_t size)
	{
		constexpr const auto len = sizeof(MemBlockHeader);
		if (!spMemBlock)
		{
			spMemBlock = std::make_shared<MemBlock>(topic, SharedMemory::ReadWrite, size + len);
		}

		if (spMemBlock->Length() >= size)
		{
			return spMemBlock->Write(data, size);
		}
		else
		{
			// no enough memory
			return ERR_IPC::E_NO_ENOUGH_MEM;
		}
	}

private:
	std::string topic;
	std::shared_ptr<MemBlock> spMemBlock{nullptr};
};

/**
 * @brief  Just read the topic
 * @author livio liu[livliu@cisco.com]
 */
class Subscriber final
{
public:
	using SubscribeCallback = std::function<void(void *, size_t)>;
	explicit Subscriber(const std::string &topic)
	{
		spMemBlock = std::make_shared<MemBlock>(topic, SharedMemory::ReadWrite);
		if (spMemBlock->IsAttached())
		{ 
			spMemBlock->IncreaseSubscriber();
		}
	}

	~Subscriber()
	{
		if (spSubscribeThread)
		{
			mAutoRead.store(false);
			spSubscribeThread->join();
			spSubscribeThread.reset();
		}
		if (spMemBlock->IsAttached())
		{
			spMemBlock->DecreaseSubscriber();
		}
	}

	bool BeginAutoSubscribe(SubscribeCallback callback)
	{
		if (!spSubscribeThread)
		{
			mCb = callback;
			mAutoRead.store(true);
			spSubscribeThread = std::make_shared<std::thread>([this]()
															  { Subscribe(); });
		}

		return spSubscribeThread != nullptr;
	}

protected:
	void Subscribe()
	{
		uint16_t curPktIdx = 0;
		char *spBuf = nullptr;
		int bufLen = 0;
		while (mAutoRead.load())
		{
			if (spMemBlock && spMemBlock->IsAttached())
			{
				if (!spBuf)
				{
					bufLen = spMemBlock->Length();
					spBuf = new char[bufLen];
				}
				auto tmpIndex = curPktIdx;
				auto err = spMemBlock->Read(spBuf, bufLen, curPktIdx);
				if (mCb && tmpIndex != curPktIdx)
				{
					mCb(spBuf, bufLen);
				}
			}
			std::this_thread::sleep_for(std::chrono::milliseconds(10));
		}

		if (spBuf)
		{
			delete[] spBuf;
			spBuf = nullptr;
		}
	}

private:
	std::shared_ptr<MemBlock> spMemBlock{nullptr};
	std::shared_ptr<std::thread> spSubscribeThread{nullptr};
	SubscribeCallback mCb{nullptr};
	std::atomic<bool> mAutoRead{false};
};

#endif // IPC2_H